package com.shen.llmtool.modules.deepseek.service;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.shen.llmtool.common.constant.CsvConstant;
import com.shen.llmtool.common.constant.LLMConstant;
import com.shen.llmtool.common.constant.PromptConstant;
import com.shen.llmtool.common.utils.CsvFileUtil;
import com.shen.llmtool.common.utils.OkHttpUtil;
import com.shen.llmtool.modules.deepseek.vo.ModelMessageVo;
import com.shen.llmtool.modules.deepseek.vo.ModelRequestVo;
import com.shen.llmtool.modules.deepseek.vo.ModelResponseVo;
import com.shen.llmtool.modules.deepseek.vo.OrganizationCategoryVo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * 机构分类
 */
@Slf4j
public class OrganizationClassifyService {
    public static final ObjectMapper mapper = new ObjectMapper();
    public static final List<String[]> writeData = new ArrayList<>();

    public static final Integer BATCH_SIZE = 35;

    public static void main(String[] args) {
        String csvReadFilePath = "E:\\city.csv";
        List<String> stringList = CsvFileUtil.readCsvFile(csvReadFilePath);
        int batchSize = 0;
        StringJoiner stringJoiner = new StringJoiner("\n");
        List<String[]> list = new ArrayList<>();
        try {
            for (String jgmc : stringList) {
                batchSize++;
                stringJoiner.add(jgmc);
                String[] record = new String[15];
                record[0] = jgmc;
                list.add(record);
                if (batchSize >= BATCH_SIZE) {
                    handleClassifyLLM(stringJoiner.toString());
                    batchSize = 0;
                    stringJoiner = new StringJoiner("\n");
                    list = new ArrayList<>();
                }
            }
        } catch (Exception e) {
            log.error("调用大模型处理过程中出现异常:", e);
        }
        CsvFileUtil.writeCsvFile("E:\\organizationClassify.csv", writeData, CsvConstant.ORGANIZATION_CSV_HEADER);
    }

    private static void handleClassifyLLM(String jgmcListStr) throws Exception {
        System.out.println(jgmcListStr);
        String jgdwlxPrompt = PromptConstant.JGDWLX_PROMPT_WORDS + jgmcListStr;
        String zzjglxPrompt = PromptConstant.ZZJGLX_PROMPT_WORDS + jgmcListStr;
        String jgcjPrompt = PromptConstant.JGCJ_PROMPT_WORDS + jgmcListStr;
        String sslyPrompt = PromptConstant.SSLY_PROMPT_WORDS + jgmcListStr;
        String bmywfcPrompt = PromptConstant.BMYWFC_PROMPT_WORDS + jgmcListStr;
        String bmzzlxPrompt = PromptConstant.BMZZLX_PROMPT_WORDS + jgmcListStr;
        String sbhPrompt = PromptConstant.SBH_PROMPT_WORDS + jgmcListStr;
        List<String> promptList = Arrays.asList(jgdwlxPrompt, zzjglxPrompt, jgcjPrompt, sslyPrompt, bmywfcPrompt, bmzzlxPrompt, sbhPrompt);
        //方式一: 异步调用快但设置过多机构数据时会出现超时
        List<List<OrganizationCategoryVo>> modelResponseList = asyncCallLLM(promptList);
        //方式二: 串行调用慢但召回率高
//        List<List<OrganizationCategoryVo>> modelResponseList = new ArrayList<>();
//        for (String prompt : promptList) {
//            modelResponseList.add(requestLLM(prompt));
//        }
        List<String[]> recordList = handleRecord(modelResponseList);
        writeData.addAll(recordList);
    }

    /**
     * 处理结果转换为一行行记录
     *
     * @param modelResponseList 模型返回结果集合
     * @return 行记录集合
     */
    private static List<String[]> handleRecord(List<List<OrganizationCategoryVo>> modelResponseList) {
        List<String[]> recordList = new ArrayList<>();
        boolean isEmptyFlag = false;
        for (List<OrganizationCategoryVo> list : modelResponseList) {
            if (CollectionUtils.isEmpty(list)) {
                isEmptyFlag = true;
            }
        }
        if (isEmptyFlag) {
            return Collections.emptyList();
        }
        //转换为二维数组
        OrganizationCategoryVo[][] arrays = new OrganizationCategoryVo[modelResponseList.size()][];
        for (int i = 0; i < modelResponseList.size(); i++) {
            List<OrganizationCategoryVo> currentList = modelResponseList.get(i);
            arrays[i] = currentList.toArray(new OrganizationCategoryVo[0]);
        }
        for (int i = 0; i < arrays[0].length; i++) {
            String[] record = new String[15];
            for (int j = 0; j < arrays.length; j++) {
                OrganizationCategoryVo organizationCategoryVo = arrays[j][i];
                if (j == 0) {
                    record[j] = organizationCategoryVo.getOrgName();
                    record[j + 1] = organizationCategoryVo.getOrgCategoryNo();
                    record[j + 2] = organizationCategoryVo.getOrgCategoryName();
                } else {
                    record[2 * j + 1] = organizationCategoryVo.getOrgCategoryNo();
                    record[2 * j + 2] = organizationCategoryVo.getOrgCategoryName();
                }
            }
            recordList.add(record);
        }
        return recordList;
    }

    /**
     * 异步调用大模型
     *
     * @param promptList 提示词集合
     * @return 大模型调用结果
     */
    private static List<List<OrganizationCategoryVo>> asyncCallLLM(List<String> promptList) {
        ExecutorService executor = Executors.newFixedThreadPool(8);
        List<List<OrganizationCategoryVo>> resultList = new ArrayList<>();
        try {
            List<Future<List<OrganizationCategoryVo>>> futureList = new ArrayList<>();
            for (String prompt : promptList) {
                Future<List<OrganizationCategoryVo>> future = executor.submit(new Callable<List<OrganizationCategoryVo>>() {
                    @Override
                    public List<OrganizationCategoryVo> call() throws Exception {
                        return requestLLM(prompt);
                    }
                });
                futureList.add(future);
            }

            for (Future<List<OrganizationCategoryVo>> future : futureList) {
                try {
                    resultList.add(future.get());
                } catch (Exception e) {
                    log.warn("执行任务异常");
                }
            }
        } finally {
            executor.shutdown();
        }
        return resultList;
    }

    /**
     * 请求大模型
     *
     * @param prompt 提示词
     * @return 组织分类对象集合
     * @throws Exception 异常
     */
    private static List<OrganizationCategoryVo> requestLLM(String prompt) throws Exception {
        Map<String, String> headersMap = buildRequestHeader();
        String modelRequestVoStr = buildRequestStr(prompt);
        String responseStr = OkHttpUtil.post(headersMap, LLMConstant.DEEPSEEK_BASE_URL, modelRequestVoStr);
        if (!StringUtils.hasText(responseStr)) {
            log.warn("[请求大模型]: 请求大模型异常");
            return Collections.emptyList();
        }
        ModelResponseVo modelResponseVo = mapper.readValue(responseStr, new TypeReference<ModelResponseVo>() {
        });
        if (isModelResponseVoEmpty(modelResponseVo)) {
            log.warn("[请求大模型]: 请求大模型响应值异常");
            return Collections.emptyList();
        }
        String formatStr = matchStr(modelResponseVo.getChoices().get(0).getMessage().getContent());
        return mapper.readValue(formatStr, new TypeReference<List<OrganizationCategoryVo>>() {
        });
    }

    /**
     * 对大模型返回内容进行格式化
     *
     * @param responseStr 大模型返回内容
     * @return 格式化后文本内容
     */
    private static String matchStr(String responseStr) {
        if (!StringUtils.hasText(responseStr)) {
            log.warn("[格式化返回值]: 返回值字符串为空");
            return "";
        }
        String formatFirst = responseStr.replaceAll("[\\n\\s\\t\\r`]", "");
        return formatFirst.replace("json", "");
    }

    /**
     * 构建LLM请求头
     *
     * @return 请求头集合
     */
    private static Map<String, String> buildRequestHeader() {
        Map<String, String> headersMap = new HashMap<>(1 << 2);
        headersMap.put(LLMConstant.AUTHORIZATION, LLMConstant.BEARER + LLMConstant.DASHSCOPE_API_KEY);
        headersMap.put(LLMConstant.CONTENT_TYPE, LLMConstant.JSON_CONTENT_TYPE);
        return headersMap;
    }

    /**
     * 构建LLM请求体字符串
     *
     * @param prompt 提示词
     * @return 请求体字符串
     * @throws Exception 异常
     */
    private static String buildRequestStr(String prompt) throws Exception {
        ModelMessageVo modelMessageVo = new ModelMessageVo("user", prompt);
        List<ModelMessageVo> modelMessageVoList = new ArrayList<>();
        modelMessageVoList.add(modelMessageVo);
        ModelRequestVo modelRequestVo = new ModelRequestVo(LLMConstant.DEEPSEEK_R1_DISTILL_70B_MODEL, modelMessageVoList);
        return mapper.writeValueAsString(modelRequestVo);
    }

    /**
     * 判断模型调用结果是否为空
     *
     * @param modelResponseVo 模型返回值
     * @return 是否为空
     */
    private static boolean isModelResponseVoEmpty(ModelResponseVo modelResponseVo) {
        return (modelResponseVo == null) || (CollectionUtils.isEmpty(modelResponseVo.getChoices()) || (modelResponseVo.getChoices().get(0) == null)
                || (modelResponseVo.getChoices().get(0).getMessage() == null) || !(StringUtils.hasText(modelResponseVo.getChoices().get(0).getMessage().getContent())));
    }


}
